#
# messkit.netstring
#
# Copyright 2006-2008 Helsinki Institute for Information Technology (HIIT)
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""
Self-delimiting string encoding proposed by D. J. Bernstein

See http://cr.yp.to/proto/netstrings.txt .

A string of 8-bit bytes is encoded as [len]":"[string]",".
For example, "hello world!" is encoded as "12:hello world!,".
"""

def make_netstring(content):
    """
    Return the content wrapped in a netstring.
    """
    content = str(content)
    return '%d:%s,' % (len(content), content)

def parse_netstring(string, start=0):
    """
    Parse a netstring starting at the specified index.

    The return value is a tuple (s, k), where s is the netstring's
    content and k is the total number of bytes in the undecoded
    netstring.  If the end of the input string is reached before the
    end of the netstring is seen, (None, k) is returned.  If the end
    of the input string is reached before a colon is seen, (None, 0)
    is returned.  If the netstring or nonempty fragment thereof is
    invalid, ValueError is raised.
    """
    try:
        colon = string.index(':', start)
    except ValueError:
        if start < len(string) and not string[start:].isdigit():
            raise ValueError, 'Invalid length in netstring fragment.'
        return None, 0

    head = string[start:colon]

    try:
        n = int(head)
    except ValueError:
        raise ValueError, 'Invalid length in netstring.'

    if str(n) != head:
        raise ValueError, 'Invalid length in netstring.'

    if n < 0:
        raise ValueError, 'Negative length in netstring.'

    end = colon + n + 2
    if len(string) < end:
        return None, end - start

    if string[end - 1] != ',':
        raise ValueError, 'Expected "," at end of netstring.'

    return string[colon + 1: end - 1], end - start

def split_netstrings(string, start=0):
    """
    Parse zero or more concatenated netstrings.

    The return value is a list of the contents of the netstrings,
    which must use up the entire input string starting at the
    specified index.
    """
    result = []
    while start < len(string):
        s, k = parse_netstring(string, start)
        if s is None:
            raise ValueError, 'Input string ended abruptly.'

        result.append(s)
        start += k

    return result

if __name__ == '__main__':
    def check_value_error(f, *args):
        try:
            f(*args)
            assert False, 'Expected ValueError.'
        except ValueError:
            pass

    assert make_netstring('') == '0:,'
    assert make_netstring('a') == '1:a,'
    assert make_netstring('abc') == '3:abc,'

    assert parse_netstring('0:,') == ('', 3)
    assert parse_netstring('1:a,') == ('a', 4)
    assert parse_netstring('3:abc,') == ('abc', 6)

    assert parse_netstring('0:,xyz') == ('', 3)
    assert parse_netstring('1:a,xyz') == ('a', 4)
    assert parse_netstring('3:abc,xyz') == ('abc', 6)

    assert parse_netstring('xyz0:,', 3) == ('', 3)
    assert parse_netstring('xyz1:a,', 3) == ('a', 4)
    assert parse_netstring('xyz3:abc,', 3) == ('abc', 6)

    assert parse_netstring('') == (None, 0)
    assert parse_netstring('0') == (None, 0)
    assert parse_netstring('0:') == (None, 3)
    assert parse_netstring('1:a') == (None, 4)
    assert parse_netstring('3:abc') == (None, 6)

    check_value_error(parse_netstring, ':')
    check_value_error(parse_netstring, 'x:')
    check_value_error(parse_netstring, '-1:')

    check_value_error(parse_netstring, '03:abc')
    check_value_error(parse_netstring, ' 3:abc')
    check_value_error(parse_netstring, '3 :abc')

    check_value_error(parse_netstring, '0x')
    check_value_error(parse_netstring, '0:x')
    check_value_error(parse_netstring, '1:ax')
    check_value_error(parse_netstring, '3:abcx')

    assert split_netstrings('') == []
    assert split_netstrings('0:,') == ['']
    assert split_netstrings('0:,1:a,3:abc,') == ['', 'a', 'abc']
    assert split_netstrings('xyz0:,1:a,3:abc,', 3) == ['', 'a', 'abc']
    check_value_error(split_netstrings, '0:,1:a,3:abc')

    print 'OK'
